import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
# 线性回归
def cal_w_b(x, y):
    x_mean = x.sum()/x.size
    fen_mu = 0.0
    xi2 = 0.0
    b = 0.0
    for xi, yi in zip(x, y):
        fen_mu += yi*(xi-x_mean)
        xi2 += xi**2
    fen_zi = xi2 - (x.sum())**2/x.size
    w = fen_mu/fen_zi
    for xi, yi in zip(x, y): 
        b += (yi-w*xi)
    return w, b/x.size

# 预测函数
def y(w, b, x):
    df_predict = pd.DataFrame(x, columns=['面积(平方)'])
    l = []
    for i in x:
        l.append(w*i+b)
    df_predict['预测价格(万)'] = l
    return df_predict

#  数据集
df_data = [
    [77.36, 470],
    [116.74, 730],
    [116.7, 760],
    [100.68, 680],
    [116.1, 700],
    [115.81, 720],
    [104.24, 700],
    [106.73, 690],
    [115.86, 730]
]
df_house = pd.DataFrame(df_data, columns=['面积(平方)', '价格(万)'])

# 测试数据
test_date = [56.6, 78.4, 58, 123.5, 56.8, 77, 150.6]

# 预测
w, b = cal_w_b(df_house['面积(平方)'], df_house['价格(万)'])
print('w:', w, 'b:', b)
df_predict = y(w, b, test_date)
print(df_predict)

# 可视化
# 显示中文
plt.rcParams['font.sans-serif'] = 'SimHei' 
plt.rcParams['axes.unicode_minus'] = False

plt.xlabel('面积(平方)')
plt.ylabel('价格(万)')
plt.scatter(df_house['面积(平方)'], df_house['价格(万)'])
plt.plot(df_predict['面积(平方)'], df_predict['预测价格(万)'])
plt.legend(['数据集','预测'])
plt.savefig('线性回归.png')